
args=(commandArgs(TRUE))

#evaluate the arguments
# input arguments will be lambda (lb), dispersal cost (dc), file.ID (fid), replicate runs (rr)
for(i in 1:length(args)) {
	 eval(parse(text=args[[i]]))
}


dyn.load("/home1/30/jc227089/evo-dispersal/LD/Functions.so")



#### Postprocessing functions ####

#calculates linkage disequilibrium scaled against Dmax
# assumes data are genotypes in a matrix [individual][D1a, D1b,...Nna, Nnb]
# assumes diploid, bi-allelic system
# pass both dispersal and neutral arrays to mtrx 
# returns mean R = D/Dmax across all locus pairs aggregated by pair type
ld<-function(mtrx, nloci, ndisploci, scale=TRUE){
	gtype.pairs<-matrix(combn(1:nloci, 2), nrow=2) #find pairwise combinations of loci
	a.pairs<-matrix(seq(1, 2*nloci, by=2)[gtype.pairs], nrow=2) #a chromosome column pairs
	b.pairs<-matrix(seq(2, 2*nloci, by=2)[gtype.pairs], nrow=2) #b chromosome column pairs
	R<-c() #vector to take scaled D values
	p.type<-c() #vector to take pair type ("DD", "NN" or "DN")
	loc1<-c()
	loc2<-c()
	for (ii in 1:length(gtype.pairs[1,])){ #step through locus pairs
		L1<-factor(mtrx[,a.pairs[1,ii]], levels=0:1)
		L2<-factor(mtrx[,a.pairs[2,ii]], levels=0:1)
		gam.fqs<-table(L1, L2) #creates table of P11, P12, etc.
		L1<-factor(mtrx[,b.pairs[1,ii]], levels=0:1)
		L2<-factor(mtrx[,b.pairs[2,ii]], levels=0:1)
		gam.fqs<-gam.fqs+table(L1, L2) #creates table of P11, P12, etc.
		gam.fqs<-gam.fqs/sum(gam.fqs) #as relative frequencies
		p1p2<-apply(gam.fqs, 1, sum) #vectors of allele frequencies
		q1q2<-apply(gam.fqs, 2, sum)
		if (sum(c(p1p2, q1q2)>0.9)>=1) next #avoid stochastic mess
		D<-gam.fqs[1,1]-p1p2[1]*q1q2[1]
		if (D>0 && scale==T) { #scale
			Dmax<-min(p1p2[1]*q1q2[2], p1p2[2]*q1q2[1])
			D<-D/Dmax	
		}
		if (D<0 && scale==T) { #scale
			Dmax<-max(-p1p2[1]*q1q2[1], -p1p2[2]*q1q2[2])
			D<-D/Dmax	
		}
		R<-c(R, D)
		pt<- "DN"
		if (gtype.pairs[1,ii]<=ndisploci & gtype.pairs[2,ii]<=ndisploci) pt<-"DD"
		if (gtype.pairs[1,ii]>ndisploci & gtype.pairs[2,ii]>ndisploci) pt<-"NN"
		id<-gtype.pairs[1,ii]
		loc1<-c(loc1, id)
		id<-gtype.pairs[2,ii]
		loc2<-c(loc2, id)
		p.type<-c(p.type, pt)
	}
	R<-data.frame(p.type, loc1, loc2, R)
	return(R)		
}




#runs a transect along the invasion front, 
# calculates LD at D and N loci in a sample of a given neigh.size along the front
front.tsect<-function(pop.mtrx, dens.mtrx, neigh.size, spY, nloci, ndisploci){
	pop.mtrx<-subset(pop.mtrx, pop.mtrx[,1]!=-9) #remove -9s
	pop.mtrx[,1] = pop.mtrx[,1]+1 #move grid notation into R format
	pop.mtrx[,2] = pop.mtrx[,2]+1
	for (yy in 1:spY){ #trim matrix to only values within neigh.size of front
		loc.max.x<-max(cumsum(dens.mtrx[yy,]>10)) #maximum x per y where grid cell population >10 individuals
		pop.mtrx<-subset(pop.mtrx, pop.mtrx[,2]!=yy | (pop.mtrx[,2]==yy & pop.mtrx[,1]>(loc.max.x-neigh.size)))
	}
	tsect.sites<-rep(1:(spY%/%neigh.size), each=neigh.size) #factors for y
	pop.mtrx<-subset(pop.mtrx, pop.mtrx[,2]<=length(tsect.sites)) #remove rows falling outside factorisation
	pop.mtrx[,2]<-tsect.sites[pop.mtrx[,2]] #convert y to factor levels
	all.mtrx<-split(pop.mtrx[,-(1:2)], pop.mtrx[,2]) #split data into list
	all.mtrx<-lapply(all.mtrx, matrix, ncol=nloci*2)
	all.mtrx<-lapply(all.mtrx, ld, nloci=nloci, ndisploci=ndisploci)
	all.mtrx<-do.call("rbind", all.mtrx)	
	all.mtrx
}


#calculates cline in allele frequencies across X
freq.cline<-function(pop.mtrx, nloci){
	pop.mtrx<-subset(pop.mtrx, pop.mtrx[,1]!=-9) #remove -9s
	n.inds<-tapply(pop.mtrx[,1], pop.mtrx[,1], length)
	n.inds<-matrix(rep(n.inds, nloci), ncol=nloci)
	agg<-rep(1:nloci, each=2)
	freq1<-aggregate(pop.mtrx[,-(1:2)], by=list(pop.mtrx[,1]), FUN=sum)
	freq1<-t(aggregate(t(freq1[,-1]), by=list(agg), FUN=sum)[,-1])
	freq1<-freq1/(2*n.inds)
	freq1
}

#calculates Nei's genetic distance between two sites at all loci
# takes matrices of individuals minus locality
neis.dist<-function(mtrx1, mtrx2, nloci){
	agg<-rep(1:nloci, each=2)
	freq1<-colSums(mtrx1) # find p for each locus
	freq1<-tapply(freq1, agg, sum)/(2*nrow(mtrx1))
	freq2<-colSums(mtrx2)
	freq2<-tapply(freq2, agg, sum)/(2*nrow(mtrx2))
	Jxy<-freq1*freq2+(1-freq1)*(1-freq2)
	Jxx<-freq1^2+(1-freq1)^2
	Jyy<-freq2^2+(1-freq2)^2
	norm.I<-Jxy/sqrt(Jxx*Jyy)
	Ns.dist<--log(norm.I)
	Ns.dist	
}

# calculates pairwise genetic distance between 
 # five populations sampled through invasion history
 #returns ls regression coefficient for I by D for each locus
dist.tsect<-function(pop.mtrx, dens.mtrx, nloci, neigh.size, ndists, xmin=10, xmax=-9){
	if (xmax==-9) xmax<-sum(dens.mtrx[1,]>10)-neigh.size
	samp.locs<-floor(seq(xmin, xmax, length.out=ndists))
	combs<-matrix(combn(samp.locs, 2), nrow=2)
	x.inds<-c()
	for (ii in 1:5){
		x.inds<-c(x.inds, samp.locs[ii]:(samp.locs[ii]+neigh.size-1))
	}
	pop.mtrx<-subset(pop.mtrx, pop.mtrx[,1]%in%x.inds & pop.mtrx[,2]<neigh.size)
	samp.locs<-rep(samp.locs, each=neigh.size)
	pop.mtrx[,1]<-samp.locs[as.numeric(as.factor(pop.mtrx[,1]))]
	pop.mtrx<-split(pop.mtrx[,-(1:2)], pop.mtrx[,1])
	pop.mtrx<-lapply(pop.mtrx, matrix, ncol=2*nloci)
	dist<-abs(combs[1,]-combs[2,])
	combs<-matrix(combn(1:ndists, 2), nrow=2)
	pop.prs<-c()
	for (ii in 1:ncol(combs)){
		pr<-neis.dist(pop.mtrx[[combs[1,ii]]], pop.mtrx[[combs[2,ii]]], nloci)
		pr<-c(dist[ii], pr)
		pop.prs<-rbind(pop.prs, pr)	
	}
	b<-c()
	for (ii in 2:ncol(pop.prs)){
		b<-c(b, lm(pop.prs[,ii]~pop.prs[,1])$coefficients[2])
	}
	matrix(b, nrow=1)	
}

#trimmed version of dist.tsect same but passed a matrix with a range of x and a single y
ibd<-function(trim.mtrx, nloci){
	combs<-matrix(combn(min(trim.mtrx[,1]):max(trim.mtrx[,1]),2), nrow=2)
	pop.mtrx<-split(trim.mtrx[,-(1:2)], trim.mtrx[,1])
	pop.mtrx<-lapply(pop.mtrx, matrix, ncol=2*nloci)
	dist<-abs(combs[1,]-combs[2,])
	combs<-combs-min(combs)+1
	pop.prs<-c()
	for (ii in 1:ncol(combs)){
		pr<-neis.dist(pop.mtrx[[combs[1,ii]]], pop.mtrx[[combs[2,ii]]], nloci)
		pr<-c(dist[ii], pr)
		pop.prs<-rbind(pop.prs, pr)	
	}
	b<-c()
	for (ii in 2:ncol(pop.prs)){
		b<-c(b, lm(pop.prs[,ii]~pop.prs[,1])$coefficients[2])
	}
	matrix(b, nrow=1)	
	
}
# calculates the fequency of 1 alleles at each locus
f1<-function(pop.mtrx, nloci){
	pop.mtrx<-pop.mtrx[,-(1:2)]
	agg<-rep(1:nloci, each=2)
	freq<-colSums(pop.mtrx)	
	freq<-tapply(freq, agg, sum)/(2*nrow(pop.mtrx))
	matrix(freq, nrow=1)
}

#defines a moving window, and runs dist.tsect and f1 on that window
m.window<-function(pop.mtrx, dens.mtrx, wind.size, nloci){
	xmin<-10
	xmax<-sum(dens.mtrx[1,]>50)
	out<-NULL
	for (ii in xmin:(xmax-wind.size)){
		x.inds<-ii:(ii+wind.size-1)
		temp<-subset(pop.mtrx, pop.mtrx[,1]%in%x.inds & pop.mtrx[,2]==1)
		a<-ibd(temp, nloci=nloci)
		b<-f1(temp, nloci)
		c<-cbind(x=ii, ibd=a, f1=b)
		out<-rbind(out, c)		
	}	
	out
}

##### Model runs ######

spY<-10
exp.gens<-101 #generations of expansion
nloci<-10 ##check this matches compiled code
ns<-5 #neighbourhood size of sample
report.rate<-50
genseq<-seq(0, exp.gens-1, report.rate)[-1]

ld.out<-NULL
IbyD.out<-NULL
MW.out<-NULL
for (r in 1:rr){
	# R_n - starting population size
	# R_spX & R_spY - starting habitat matrix size
	# R_lambda - basic number of offspring
	# R_K - carrying capacity
	# R_disp_cost - cost of dispersing .. ranges between 0 & 1
	# R_reports_rate - rate at which outputs of density & individuals are sent back to R
	# R_num_gens - the number of generations past the burn in period... plus also the increase spX to be used
	# rho - the R working environment
	.Call('mother',R_n=10000,R_spX=10,R_spY=spY,R_lambda=lb,R_K=1000,R_disp_cost=dc,R_reports_rate=report.rate,R_num_gens=exp.gens,rho=environment(), R_m_rate=0.0001)
	
	# summarize ld across space and generations
	for (gg in 1:length(genseq)){ #step through generations
		temp<-get(paste("gen_", genseq[gg], sep=""))
		temp.dens<-get(paste("dens_", genseq[gg], sep=""))
		ac.ns<-front.tsect(temp, temp.dens, ns, spY, 2*nloci, nloci)
		ac.ns<-data.frame(dc=dc, lambda=lb, gen=genseq[gg], neigh.size=ns, rep=r, ac.ns)
		ibd<-dist.tsect(temp, temp.dens, 2*nloci, ns, 5, xmin=5, xmax=-9)
		ibd<-data.frame(dc=dc, lambda=lb, gen=genseq[gg], neigh.size=ns, rep=r, ibd)
		mw<-m.window(temp, temp.dens, 5, 2*nloci)
		mw<-data.frame(dc=dc, lambda=lb, gen=genseq[gg], neigh.size=ns, rep=r, mw)
		if (gg==1 & r==1) {
			ld.out<-ac.ns
			IbyD.out<-ibd
			MW.out<-mw
		}
		else {
			ld.out<-rbind(ld.out, ac.ns)
			IbyD.out<-rbind(IbyD.out, ibd)
			MW.out<-rbind(MW.out, mw)
		}
	}
}

save(ld.out, file=paste("/home1/30/jc227089/LD_inv/ran_outs/ldout", fid, ".RData", sep=""))
save(IbyD.out, file=paste("/home1/30/jc227089/LD_inv/ran_outs/IbyDout", fid, ".RData", sep=""))
save(MW.out, file=paste("/home1/30/jc227089/LD_inv/ran_outs/MWout", fid, ".RData", sep=""))



